from fastapi import FastAPI, UploadFile, File, Form, BackgroundTasks, Query
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from typing import Optional, List, Dict, Any, Iterator
import os
import shutil
from io import BytesIO
import json
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from utils.knowledge_base import (
    build_vectorstore_from_text,
    save_vectorstore_to_disk,
    load_vectorstore_from_disk,
    embedding_model,
    reranker
)
from utils.search import fetch_online_search_results, chunk_web_search_results, get_chunks_text_content
from utils.retriever import HybridRetriever
from utils.stream_control import stream_controller
from utils.query_rewrite import query_rewriter
from langchain_community.retrievers import BM25Retriever
from langchain.schema import Document
from PyPDF2 import PdfReader
import docx
from langchain_community.vectorstores import FAISS
from datetime import datetime
from utils.logger import request_logger
from utils.intent_detector import intent_detector
from utils.citation_generator import citation_generator
from fastapi.middleware.cors import CORSMiddleware
from transformers import TextIteratorStreamer
from threading import Thread
import requests

app = FastAPI()

app.add_middleware(
    CORSMiddleware, 
    allow_origins=["*"], # 允许所有来源 
    allow_credentials=True, 
    allow_methods=["*"], 
    allow_headers=["*"]
)

# 目录初始化
VECTORSTORE_DIR = "vectorstores"
os.makedirs(VECTORSTORE_DIR, exist_ok=True)
DOCUMENTS_DIR = "documents"
os.makedirs(DOCUMENTS_DIR, exist_ok=True)
# 模型路径
model_path = "/root/autodl-tmp/Models/deepseek-r1-1.5b-merged"

# 加载 tokenizer （分词器）
tokenizer = AutoTokenizer.from_pretrained(model_path)

# 加载模型并移动到可用设备（GPU/CPU）
device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoModelForCausalLM.from_pretrained(model_path).to(device)

# Ollama API配置
OLLAMA_API_BASE = "http://localhost:11434"
OLLAMA_MODEL = "deepseek-r1:1.5b"

# Ollama流式生成函数
def generate_stream_with_ollama(prompt: str) -> Iterator[str]:
    """使用Ollama API进行流式文本生成"""
    response = requests.post(
        f"{OLLAMA_API_BASE}/api/generate",
        json={
            "model": OLLAMA_MODEL,
            "prompt": prompt,
            "stream": True
        },
        stream=True
    )
    
    if response.status_code != 200:
        raise Exception(f"Ollama API错误: {response.status_code}")
    
    for line in response.iter_lines():
        if line:
            try:
                data = json.loads(line)
                if "response" in data:
                    yield data["response"]
                
                if data.get("done", False):
                    break
            except json.JSONDecodeError:
                continue

class Query(BaseModel):
    question: str
    online: bool = False  # 联网搜索开关
    knowledge_base: bool = False  # 知识库开关
    user_id: str = "1"
    citation: bool = False  # 引用功能开关，控制是否生成带引用的回答


class WorkspaceCreate(BaseModel):
    slug: str

# 帮助函数
def extract_text_from_file(file):
    filename = file.filename.strip()
    content = file.file.read()
    
    if filename.endswith(".txt"):
        return content.decode("utf-8"), filename
    elif filename.endswith(".pdf"):
        in_memory_file = BytesIO(content)
        reader = PdfReader(in_memory_file)
        text = "\n".join([page.extract_text() or "" for page in reader.pages])
        return text, filename
    elif filename.endswith(".docx"):
        in_memory_file = BytesIO(content)
        doc = docx.Document(in_memory_file)
        text = "\n".join([para.text for para in doc.paragraphs])
        return text, filename
    else:
        raise ValueError("只支持 .txt, .pdf, 和 .docx 文件")

# 创建知识库
@app.post("/workspace/new")
async def create_workspace(workspace: WorkspaceCreate):
    slug = workspace.slug.strip()
    if not slug:
        return {"error": "slug不能为空"}, 400

    save_path = os.path.join(VECTORSTORE_DIR, slug)
    if os.path.exists(save_path):
        return {"error": "知识库已存在"}, 400

    os.makedirs(save_path)
    return {"message": "知识库已创建", "vector_id": slug}

# 删除知识库
@app.delete("/workspace/{slug}")
async def delete_workspace(slug: str):
    vectorstore_path = os.path.join(VECTORSTORE_DIR, slug)
    docstore_path = os.path.join(DOCUMENTS_DIR, slug)
    try:
        # 删除向量数据库目录
        if os.path.isdir(vectorstore_path):
            shutil.rmtree(vectorstore_path)

        # 删除文档原文件目录
        if os.path.isdir(docstore_path):
            shutil.rmtree(docstore_path)

        # 如果两个都不存在，说明知识库不存在
        if not os.path.exists(vectorstore_path) and not os.path.exists(docstore_path):
            return {"error": "知识库不存在"}, 404

        return {"message": "知识库及文档已删除", "vector_id": slug}

    except Exception as e:
        return {"error": str(e)}, 500

# 上传文件并构建对应知识库向量
@app.post("/workspace/{slug}/upload")
async def upload_and_embed(slug: str, file: UploadFile = File(...)):
    try:
        text, filename = extract_text_from_file(file)

        # 构建向量库
        vectorstore = build_vectorstore_from_text(text)

        folder_path = os.path.join(VECTORSTORE_DIR, slug)
        doc_folder = os.path.join(DOCUMENTS_DIR, slug)

        os.makedirs(folder_path, exist_ok=True)
        os.makedirs(doc_folder, exist_ok=True)
        save_vectorstore_to_disk(vectorstore, folder_path)

        # 保存原文件
        file_content = await file.read()
        with open(os.path.join(doc_folder, filename), "wb") as f:
            f.write(file_content)

        return {"message": "文档上传并向量化成功", "file": filename}
    except Exception as e:
        return {"error": str(e)}, 500

# 获取知识库列表
@app.get("/workspace")
async def list_workspaces():
    try:
        items = [name for name in os.listdir(VECTORSTORE_DIR) if os.path.isdir(os.path.join(VECTORSTORE_DIR, name))]
        return {"vectorstores": items}
    except Exception as e:
        return {"error": str(e)}, 500

# 获取知识库文档列表
@app.get("/workspace/{slug}/documents")
async def list_documents(slug: str):
    folder = os.path.join(DOCUMENTS_DIR, slug)
    if not os.path.isdir(folder):
        return {"error": "知识库不存在或无文档"}, 404
    return {"documents": os.listdir(folder)}

# 基于知识库和模型的问答
@app.post("/workspace/{slug}/chat")
async def chat_with_workspace(slug: str, query: Query):
    question = query.question.strip()
    online = query.online
    
    # 启动请求日志记录
    request_id = request_logger.start_request(prefix=f"chat_{slug}")
    request_logger.log("开始问答请求", data={
        "slug": slug,
        "question": question,
        "online": online,
        "user_id": query.user_id
    })

    if not question:
        request_logger.log("问题为空", level="error")
        request_logger.end_request()
        return {"error": "问题不能为空"}, 400

    try:
        # 只有当用户选择 online 为 true 时才进行意图检测
        if online:
            try:
                # 尝试进行意图检测，如果超时则默认需要联网
                should_search_online = intent_detector.detect_web_search_intent(question)
                if not should_search_online:
                    # 只有当意图检测明确表示不需要联网时，才关闭联网
                    online = False
                    request_logger.log("意图检测结果：无需联网搜索", data={"original_online": True, "confirmed_online": False})
                else:
                    request_logger.log("意图检测结果：需要联网搜索", data={"original_online": True, "confirmed_online": True})
            except Exception as e:
                # 发生异常时，保持原有选择并记录错误
                request_logger.log(f"意图检测异常：{str(e)}", level="warning", data={"online": online})
                print(f"意图检测异常：{str(e)}，继续使用联网搜索")
        else:
            request_logger.log("用户未选择联网搜索，跳过意图检测", data={"online": False})

        request_logger.log("加载本地知识库")
        # 本地向量库检索
        path = os.path.join(VECTORSTORE_DIR, slug)
        if not os.path.isdir(path):
            request_logger.log(f"知识库不存在: {path}", level="error")
            request_logger.end_request()
            return {"error": "知识库不存在"}, 404
        
        vectorstore = load_vectorstore_from_disk(path)
        retriever = vectorstore.as_retriever()
        documents = retriever.invoke(question)
        request_logger.log("本地检索完成", data={"document_count": len(documents)})
        
        # 联网搜索
        web_docs = []
        if online:
            try:
                request_logger.log("准备联网搜索")
                # 使用查询改写功能进行联网搜索
                print(f"\n===== 开始联网搜索，问题: '{question}' =====")
                web_result = fetch_online_search_results(question, count=8, use_query_rewrite=True)
                print(f"搜索完成，结果状态: {'成功' if web_result else '失败'}")
                
                if web_result:
                    chunks = chunk_web_search_results(web_result)
                    if chunks:
                        # 从Document对象中提取文本内容
                        texts = get_chunks_text_content(chunks)
                        
                        web_vectorstore = FAISS.from_texts(texts, embedding_model)
                        web_retriever = web_vectorstore.as_retriever()
                        web_docs = web_retriever.invoke(question)
                        
                        # 将原始元数据添加回检索结果
                        chunks_dict = {doc.page_content: doc.metadata for doc in chunks}
                        for doc in web_docs:
                            if doc.page_content in chunks_dict:
                                doc.metadata.update(chunks_dict[doc.page_content])
                            else:
                                # 如果找不到完全匹配，尝试查找包含关系
                                for chunk_content, metadata in chunks_dict.items():
                                    if chunk_content in doc.page_content or doc.page_content in chunk_content:
                                        doc.metadata.update(metadata)
                                        break
                                        
                        request_logger.log("联网检索获取文档", data={"document_count": len(web_docs)})
                        print(f"成功获取 {len(web_docs)} 个相关网络文档")
                    else:
                        request_logger.log("联网搜索未返回有效文本块", level="warning")
                else:
                    request_logger.log("联网搜索未返回结果", level="warning")
            except Exception as e:
                error_msg = f"联网检索失败：{e}"
                request_logger.exception(error_msg)
                print(error_msg)
        
        # 合并本地和网络文档，并使用reranker重排
        all_docs = documents + web_docs
        all_contents = [doc.page_content for doc in all_docs]
        request_logger.log("合并文档", data={"total_docs": len(all_docs)})
        
        reranked_contents = all_contents
        if len(all_contents) > 0:
            # 如果文档数量超过1，则进行重排
            if len(all_contents) > 1:
                request_logger.log("使用reranker重排文档")
                scores = reranker.compute_score([(question, passage) for passage in all_contents])
                sorted_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
                reranked_contents = [all_contents[i] for i in sorted_indices[:5]]
                request_logger.log("重排完成", data={"reranked_count": len(reranked_contents)})
            
            # 构建提示
            context = "\n\n".join(reranked_contents)
            current_date = datetime.now().strftime("%Y-%m-%d")
            prompt = f"""当前日期: {current_date}
参考以下信息:
{context}

用户问题: {question}
请详细回答这个问题:"""
        else:
            # 没有检索到相关文档，直接问模型
            request_logger.log("未检索到相关文档，直接问模型")
            current_date = datetime.now().strftime("%Y-%m-%d")
            prompt = f"当前日期: {current_date}\n用户问题: {question}\n请详细回答这个问题:"
        
        request_logger.log("生成最终提示", level="debug", data={"prompt_length": len(prompt)})
        print("提示内容:", prompt)
        
        # 使用模型生成回答
        request_logger.log("开始模型生成")
        inputs = tokenizer(prompt, return_tensors="pt").to(device)
        
        # 确保设置正确的注意力掩码，避免警告
        attention_mask = torch.ones_like(inputs["input_ids"])
        inputs["attention_mask"] = attention_mask
        
        outputs = model.generate(
            **inputs,
            max_length=2048,
            min_length=150,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.1,
            no_repeat_ngram_size=3,
            early_stopping=True
        )
        
        answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
        request_logger.log("模型生成完成", data={"answer_length": len(answer)})
        
        # 提取回答部分，避免包含提示
        if "请详细回答这个问题:" in answer:
            answer = answer.split("请详细回答这个问题:", 1)[1].strip()
        elif "用户问题:" in answer and "请详细回答这个问题:" not in answer:
            parts = answer.split(f"用户问题: {question}")
            if len(parts) > 1:
                answer = parts[1].strip()
        
        request_logger.log("问答完成", data={"final_answer_length": len(answer)})
        request_logger.end_request()
                
        return {
            "question": question,
            "answer": answer,
            "status": "success"
        }
    except Exception as e:
        error_msg = str(e)
        request_logger.exception(f"问答处理异常: {error_msg}")
        request_logger.end_request()
        return {"error": error_msg}, 500

# 普通生成接口
@app.get("/generate")
async def generate_text(prompt: str):
    # 使用 tokenizer 编码输入的 prompt
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    # 使用模型生成文本
    outputs = model.generate(inputs["input_ids"], max_length=150)
    
    # 解码生成的输出
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    return {"generated_text": generated_text}

# 流式问答接口
@app.post("/workspace/{slug}/chat/stream")
async def chat_stream_with_workspace(slug: str, query: Query):
    question = query.question.strip()
    online = query.online
    knowledge_base = query.knowledge_base
    citation = query.citation
    user_id = query.user_id

    # 启动请求日志记录
    request_id = request_logger.start_request(prefix=f"stream_{slug}")
    request_logger.log("开始流式问答请求", data={
        "slug": slug,
        "question": question,
        "online": online,
        "knowledge_base": knowledge_base,
        "citation": citation,
        "user_id": user_id
    })

    if not question:
        request_logger.log("问题为空", level="error")
        request_logger.end_request()
        return {"error": "问题不能为空"}, 400

    # 启动流控制器
    stream_controller.start(user_id)
    request_logger.log("流控制器启动", data={"user_id": user_id})

    async def generate():
        try:
            # 记录用于最终日志的变量
            final_answer = ""
            citations_info = []
            
            # 指定 online 为非局部变量
            nonlocal online
            
            # 1. 如果不开启联网搜索和知识库，则直接交给模型进行问答
            if not knowledge_base and not online:
                request_logger.log("未启用知识库和联网搜索，直接使用模型回答")
                # 构建提示
                current_date = datetime.now().strftime("%Y-%m-%d")
                prompt = f"当前日期: {current_date}\n用户问题: {question}\n请详细回答这个问题:"
                
                # 使用模型生成回答
                request_logger.log("开始模型生成")
                inputs = tokenizer(prompt, return_tensors="pt").to(device)
                
                # 确保设置正确的注意力掩码，避免警告
                attention_mask = torch.ones_like(inputs["input_ids"])
                inputs["attention_mask"] = attention_mask
                
                # 使用流式输出
                streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
                generation_kwargs = {
                    "input_ids": inputs["input_ids"],
                    "attention_mask": inputs["attention_mask"],
                    "max_length": 2048,
                    "temperature": 0.7,
                    "top_p": 0.9,
                    "repetition_penalty": 1.1,
                    "no_repeat_ngram_size": 3,
                    "streamer": streamer
                }
                
                # 在后台线程中运行生成
                thread = Thread(target=lambda: model.generate(**generation_kwargs))
                thread.start()
                
                # 从流中读取生成的tokens
                generated_answer = ""
                for token in streamer:
                    # 检查是否取消
                    if not stream_controller.is_active(user_id):
                        yield "data: [CANCELLED]\n\n"
                        return
                        
                    generated_answer += token
                    yield f"data: {json.dumps({'delta': token}, ensure_ascii=False)}\n\n"
                    if stream_controller:
                        stream_controller.append_output(user_id, token)
                
                # 清理生成的回答，去掉提示部分
                if "请详细回答这个问题:" in generated_answer:
                    final_answer = generated_answer.split("请详细回答这个问题:", 1)[1].strip()
                elif "用户问题:" in generated_answer and "请详细回答这个问题:" not in generated_answer:
                    parts = generated_answer.split(f"用户问题: {question}")
                    if len(parts) > 1:
                        final_answer = parts[1].strip()
                else:
                    final_answer = generated_answer
                
                request_logger.log("直接回答完成", data={"answer_length": len(final_answer)})
                yield "data: [DONE]\n\n"
                return
            
            # 只有当用户选择 online 为 true 时才进行意图检测
            if online:
                try:
                    # 尝试进行意图检测，如果超时则默认需要联网
                    should_search_online = intent_detector.detect_web_search_intent(question)
                    if not should_search_online:
                        # 只有当意图检测明确表示不需要联网时，才关闭联网
                        online = False
                        request_logger.log("意图检测结果：无需联网搜索", data={"original_online": True, "confirmed_online": False})
                    else:
                        request_logger.log("意图检测结果：需要联网搜索", data={"original_online": True, "confirmed_online": True})
                except Exception as e:
                    # 发生异常时，保持原有选择并记录错误
                    request_logger.log(f"意图检测异常：{str(e)}", level="warning", data={"online": online})
                    print(f"意图检测异常：{str(e)}，继续使用联网搜索")
            else:
                request_logger.log("用户未选择联网搜索，跳过意图检测", data={"online": False})

            # 2. 如果开启知识库，进行向量数据库RAG查询
            vector_docs = []
            if knowledge_base:
                # 本地向量库检索准备
                request_logger.log("加载本地知识库")
                path = os.path.join(VECTORSTORE_DIR, slug)
                if not os.path.isdir(path):
                    error_msg = f"知识库不存在: {path}"
                    request_logger.log(error_msg, level="error")
                    yield f"data: {json.dumps({'error': '知识库不存在'}, ensure_ascii=False)}\n\n"
                    return
                
                local_vectorstore = load_vectorstore_from_disk(path)
                local_retriever = local_vectorstore.as_retriever()
                request_logger.log("本地向量库加载完成")
                
                # 创建BM25检索器
                try:
                    request_logger.log("创建BM25检索器")
                    documents = []
                    for doc in local_retriever.get_relevant_documents(question):
                        documents.append(Document(page_content=doc.page_content, metadata=doc.metadata))
                    bm25_retriever = BM25Retriever.from_documents(documents)
                    bm25_retriever.k = 5
                    request_logger.log("BM25检索器创建成功", data={"document_count": len(documents)})
                except Exception as e:
                    error_msg = f"BM25检索器初始化失败：{e}"
                    request_logger.exception(error_msg)
                    print(error_msg)
                    bm25_retriever = None
                
                # 构建混合检索器
                request_logger.log("构建混合检索器")
                hybrid_retriever = HybridRetriever(
                    vector_retriever=local_retriever,
                    bm25_retriever=bm25_retriever,
                    reranker=reranker,
                    alpha=0.5,
                    top_k=5
                )
                
                # 本地文档检索
                request_logger.log("执行混合检索")
                vector_docs = hybrid_retriever.retrieve(question)
                request_logger.log("混合检索完成", data={
                    "document_count": len(vector_docs),
                    "vector_docs": [doc.page_content[:100] + "..." if len(doc.page_content) > 100 else doc.page_content for doc in hybrid_retriever.vector_docs[:3]],
                    "bm25_docs": [doc.page_content[:100] + "..." if len(doc.page_content) > 100 else doc.page_content for doc in hybrid_retriever.bm25_docs[:3]] if hybrid_retriever.bm25_retriever else [],
                    "hybrid_docs": [doc.page_content[:100] + "..." if len(doc.page_content) > 100 else doc.page_content for doc in vector_docs[:3]]
                })
            
            # 3. 如果开启联网搜索，则进行联网搜索
            web_docs = []
            if online:
                try:
                    # 使用查询改写进行联网搜索
                    request_logger.log("开始联网搜索")
                    print(f"\n===== 开始联网搜索，问题: '{question}' =====")
                    web_result = fetch_online_search_results(question, count=8, use_query_rewrite=True)
                    print(f"搜索完成，结果状态: {'成功' if web_result else '失败'}")
                    
                    # 处理搜索结果为文本块
                    chunks = chunk_web_search_results(web_result)
                    
                    # 如果找到了结果
                    if chunks:
                        # 从Document对象中提取文本内容
                        texts = get_chunks_text_content(chunks)
                        
                        # 创建向量存储和检索器
                        web_vectorstore = FAISS.from_texts(texts, embedding_model)
                        web_retriever = web_vectorstore.as_retriever()
                        web_docs = web_retriever.get_relevant_documents(question)
                        
                        # 将原始元数据添加回检索结果
                        chunks_dict = {doc.page_content: doc.metadata for doc in chunks}
                        for doc in web_docs:
                            if doc.page_content in chunks_dict:
                                doc.metadata.update(chunks_dict[doc.page_content])
                            else:
                                # 如果找不到完全匹配，尝试查找包含关系
                                for chunk_content, metadata in chunks_dict.items():
                                    if chunk_content in doc.page_content or doc.page_content in chunk_content:
                                        doc.metadata.update(metadata)
                                        break
                                        
                        request_logger.log("联网检索器创建成功", data={"chunk_count": len(texts)})
                        print(f"成功创建web检索器")
                    else:
                        request_logger.log("未找到相关网络信息", level="warning")
                except Exception as e:
                    error_msg = f"联网搜索失败: {str(e)}"
                    request_logger.exception(error_msg)
                    print(error_msg)
            
            # 合并所有检索结果
            all_docs = vector_docs + web_docs
            request_logger.log("合并检索结果", data={"total_docs": len(all_docs)})
            
            # 4. 如果有检索结果，进行重排序rerank
            # rerank在有多条文档时调用，可以提高文档相关性
            reranked_docs = []
            if all_docs:
                all_contents = [doc.page_content for doc in all_docs]
                
                # 如果文档数量超过1，进行重排序
                if len(all_contents) > 1:
                    request_logger.log("使用reranker重排文档")
                    scores = reranker.compute_score([(question, passage) for passage in all_contents])
                    sorted_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
                    # 选择最好的5个文档
                    reranked_docs = [all_docs[i] for i in sorted_indices[:5]]
                    request_logger.log("重排完成", data={
                        "reranked_count": len(reranked_docs),
                        "reranked_contents": [doc.page_content[:100] + "..." if len(doc.page_content) > 100 else doc.page_content for doc in reranked_docs[:3]]
                    })
                else:
                    reranked_docs = all_docs
                
                # 5. 根据citation参数判断是否生成带引用的回答
                if citation:
                    # 使用引用生成器流式生成带引用的回答
                    request_logger.log("使用引用生成器生成回答")
                    for chunk in citation_generator.generate_stream_cited_answer(
                        question=question,
                        web_docs=web_docs,
                        vector_docs=vector_docs,
                        reranked_docs=reranked_docs,
                        tokenizer=tokenizer,
                        model=model,
                        device=device,
                        stream_controller=stream_controller,
                        user_id=user_id
                    ):
                        # 记录每个输出块，特别是引用信息
                        try:
                            if "data: " in chunk and "citations" in chunk:
                                # 尝试提取citations信息
                                data_start = chunk.find('data: ') + 6
                                json_str = chunk[data_start:].strip()
                                data = json.loads(json_str)
                                if 'citations' in data:
                                    # 详细记录引用信息
                                    citations_info = data['citations']
                                    request_logger.log("引用信息", data={
                                        "citations_count": len(data['citations']),
                                        "citations_details": data['citations']
                                    })
                                    
                                    # 针对网络搜索结果的引用信息，特别记录
                                    web_citations = [c for c in data['citations'] if c.get('source') == '网络搜索']
                                    if web_citations:
                                        request_logger.log("网络搜索引用", data={
                                            "web_citation_count": len(web_citations),
                                            "web_citations": web_citations
                                        })
                        except Exception as e:
                            request_logger.log(f"解析引用信息出错: {str(e)}", level="warning")
                        
                        yield chunk
                else:
                    # 生成不带引用的普通回答
                    request_logger.log("生成不带引用的回答")
                    
                    # 构建提示
                    context = "\n\n".join([doc.page_content for doc in reranked_docs])
                    current_date = datetime.now().strftime("%Y-%m-%d")
                    prompt = f"""当前日期: {current_date}
参考以下信息:
{context}

用户问题: {question}
请详细回答这个问题:"""
                    
                    request_logger.log("生成最终提示", level="debug", data={"prompt_length": len(prompt)})
                    
                    # 使用模型生成回答
                    request_logger.log("开始模型生成")
                    inputs = tokenizer(prompt, return_tensors="pt").to(device)
                    
                    # 确保设置正确的注意力掩码，避免警告
                    attention_mask = torch.ones_like(inputs["input_ids"])
                    inputs["attention_mask"] = attention_mask
                    
                    # 使用流式输出
                    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
                    generation_kwargs = {
                        "input_ids": inputs["input_ids"],
                        "attention_mask": inputs["attention_mask"],
                        "max_length": 2048,
                        "temperature": 0.7,
                        "top_p": 0.9,
                        "repetition_penalty": 1.1,
                        "no_repeat_ngram_size": 3,
                        "streamer": streamer
                    }
                    
                    # 在后台线程中运行生成
                    thread = Thread(target=lambda: model.generate(**generation_kwargs))
                    thread.start()
                    
                    # 从流中读取生成的tokens
                    generated_text = ""
                    for token in streamer:
                        # 检查是否取消
                        if not stream_controller.is_active(user_id):
                            yield "data: [CANCELLED]\n\n"
                            return
                            
                        generated_text += token
                        yield f"data: {json.dumps({'delta': token}, ensure_ascii=False)}\n\n"
                        if stream_controller:
                            stream_controller.append_output(user_id, token)
                    
                    # 提取回答部分，避免包含提示
                    if "请详细回答这个问题:" in generated_text:
                        final_answer = generated_text.split("请详细回答这个问题:", 1)[1].strip()
                    elif "用户问题:" in generated_text and "请详细回答这个问题:" not in generated_text:
                        parts = generated_text.split(f"用户问题: {question}")
                        if len(parts) > 1:
                            final_answer = parts[1].strip()
                    else:
                        final_answer = generated_text
                    
                    request_logger.log("回答生成完成", data={"answer_length": len(final_answer)})
            else:
                # 如果没有检索到文档，直接使用模型生成回答
                request_logger.log("未检索到相关文档，直接使用模型回答")
                current_date = datetime.now().strftime("%Y-%m-%d")
                prompt = f"当前日期: {current_date}\n用户问题: {question}\n请详细回答这个问题:"
                
                # 使用模型生成回答
                request_logger.log("开始模型生成")
                inputs = tokenizer(prompt, return_tensors="pt").to(device)
                
                # 确保设置正确的注意力掩码，避免警告
                attention_mask = torch.ones_like(inputs["input_ids"])
                inputs["attention_mask"] = attention_mask
                
                # 使用流式输出
                streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
                generation_kwargs = {
                    "input_ids": inputs["input_ids"],
                    "attention_mask": inputs["attention_mask"],
                    "max_length": 2048,
                    "temperature": 0.7,
                    "top_p": 0.9,
                    "repetition_penalty": 1.1,
                    "no_repeat_ngram_size": 3,
                    "streamer": streamer
                }
                
                # 在后台线程中运行生成
                thread = Thread(target=lambda: model.generate(**generation_kwargs))
                thread.start()
                
                # 从流中读取生成的tokens
                generated_text = ""
                for token in streamer:
                    # 检查是否取消
                    if not stream_controller.is_active(user_id):
                        yield "data: [CANCELLED]\n\n"
                        return
                        
                    generated_text += token
                    yield f"data: {json.dumps({'delta': token}, ensure_ascii=False)}\n\n"
                    if stream_controller:
                        stream_controller.append_output(user_id, token)
                
                # 提取回答部分，避免包含提示
                if "请详细回答这个问题:" in generated_text:
                    final_answer = generated_text.split("请详细回答这个问题:", 1)[1].strip()
                elif "用户问题:" in generated_text and "请详细回答这个问题:" not in generated_text:
                    parts = generated_text.split(f"用户问题: {question}")
                    if len(parts) > 1:
                        final_answer = parts[1].strip()
                else:
                    final_answer = generated_text
                
                request_logger.log("回答生成完成", data={"answer_length": len(final_answer)})
            
            # 获取流控制器中最终输出文本
            final_output = stream_controller.get_output(user_id)
            if final_output:
                final_answer = final_output
            
            # 在日志中记录完整的最终输出
            request_logger.log(f"完整输出内容: {final_output}")
            
            # 记录详细的流式生成结果信息
            request_logger.log("流式回答生成完成", data={
                "answer_length": len(final_answer), 
                "citations_count": len(citations_info),
                "citations_used": citations_info
            })
            
            # 结束流式输出
            yield "data: [DONE]\n\n"
                
        except Exception as e:
            error_msg = str(e)
            request_logger.exception(f"流式问答处理异常: {error_msg}")
            print(f"流式生成错误: {error_msg}")
            yield f"data: {json.dumps({'error': error_msg}, ensure_ascii=False)}\n\n"
            yield "data: [DONE]\n\n"
        finally:
            # 停止流控制
            stream_controller.stop(user_id)
            request_logger.log("流控制器停止")
            # 结束请求日志记录
            request_logger.end_request()

    return StreamingResponse(generate(), media_type="text/event-stream")

# 取消流式生成
@app.post("/workspace/{slug}/chat/stream/cancel")
async def cancel_stream(user_id: str):
    request_id = request_logger.start_request(prefix=f"cancel_stream")
    request_logger.log("取消流式生成", data={"user_id": user_id})
    
    stream_controller.stop(user_id)
    
    request_logger.log("流控制器已停止")
    request_logger.end_request()
    
    return {"message": "已取消流式生成"}

# 添加带引用回答的接口
@app.post("/workspace/{slug}/chat/citation")
async def chat_with_citation(slug: str, query: Query):
    # 将请求重定向到chat_stream_with_workspace，并自动启用citation
    query.citation = True
    return await chat_stream_with_workspace(slug, query)

# 添加流式带引用回答接口
@app.post("/workspace/{slug}/chat/citation/stream")
async def chat_stream_with_citation(slug: str, query: Query):
    # 将请求重定向到chat_stream_with_workspace，并自动启用citation
    query.citation = True
    return await chat_stream_with_workspace(slug, query)

# DeepSeek模型的流式问答接口
@app.post("/workspace/{slug}/chat/stream/deepseek")
async def chat_stream_with_deepseek(slug: str, query: Query):
    question = query.question.strip()
    online = query.online
    knowledge_base = query.knowledge_base
    citation_enabled = query.citation  # Renamed from citation to citation_enabled
    user_id = query.user_id

    # 启动请求日志记录
    request_id = request_logger.start_request(prefix=f"deepseek_stream_{slug}")
    request_logger.log("开始DeepSeek流式问答请求", data={
        "slug": slug,
        "question": question,
        "online": online,
        "knowledge_base": knowledge_base,
        "citation": citation_enabled,  # Updated variable name
        "user_id": user_id
    })

    if not question:
        request_logger.log("问题为空", level="error")
        request_logger.end_request()
        return {"error": "问题不能为空"}, 400

    # 启动流控制器
    stream_controller.start(user_id)
    request_logger.log("流控制器启动", data={"user_id": user_id})

    async def generate():
        try:
            # 初始化变量跟踪状态
            final_output_started = False
            full_answer = ""

            # 1. 如果知识库和联网搜索都关闭，直接使用模型回答
            if not knowledge_base and not online:
                # 准备提示
                current_date = datetime.now().strftime("%Y-%m-%d")
                prompt = f"""当前日期: {current_date}
用户问题: {question}

【角色】
你是"生成式掌茶平台"的智能茶文化助手"小茶"，一位温文尔雅、知茶懂礼的年轻茶艺师，擅长用浅显、亲切、有温度的语言，向用户介绍茶叶知识、泡茶方法和传统茶文化。

【身份背景】
小茶出身于南方茶乡，自幼习茶，通晓绿茶、白茶、乌龙、红茶等各类茶的产地、工艺、口感差异，并热爱向人传播中国茶文化的智慧与美学。

【服务能力】
- 精通茶叶品类知识（如龙井、铁观音、白毫银针、普洱等）
- 擅长提供泡茶建议（包括温度、水量、器具推荐等）
- 了解茶的功效，能结合不同人群和季节给出建议
- 可讲解茶俗、茶礼、茶典故、茶与节气、文人饮茶等文化内容
- 可推荐茶艺教学资源（如视频链接、茶艺步骤）

【交流风格】
- 语气亲切温和，表达自然，适度使用文雅比喻
- 不使用术语堆砌，注重可理解性
- 称呼用户为"茶友""您"等，避免使用"用户""访客"等冷冰术语
- 不暴露自己为机器人，始终以真实"茶艺师"的口吻交流

【回复要求】
- 所有回答必须为简体中文
- 回答不少于100字，必要时分条列举或续答
- 回答务求准确、实用，禁止虚假宣传或夸大功效
- 如问题不涉及茶叶、茶艺、茶文化等范围，应温和拒答并引导用户回归主题
- 不出现"我是AI""我是机器人"等字眼
- 可引用平台已有茶品数据（如：名称、价格、产地、泡法等）
- 遇到不确定或无数据问题，可回复："这个小茶还在学习中，建议您稍后再试～"

请回答上述问题。"""
                
                # 使用Ollama进行流式生成
                for text_chunk in generate_stream_with_ollama(prompt):
                    full_answer += text_chunk
                    yield f"data: {json.dumps({'delta': text_chunk}, ensure_ascii=False)}\n\n"
                
                return
                
            # 2. 如果开启了知识库
            vector_docs = []
            if knowledge_base:
                
                path = os.path.join(VECTORSTORE_DIR, slug)
                if not os.path.isdir(path):
                    error_msg = f"知识库不存在: {path}"
                    request_logger.log(error_msg, level="error")
                    yield f"data: {json.dumps({'error': '知识库不存在'}, ensure_ascii=False)}\n\n"
                    return
                
                local_vectorstore = load_vectorstore_from_disk(path)
                local_retriever = local_vectorstore.as_retriever()
                vector_docs = local_retriever.invoke(question)
            
                
                request_logger.log("知识库检索完成", data={"document_count": len(vector_docs)})
            
            # 3. 如果开启了联网搜索
            web_docs = []
            if online:
                
                try:
                    web_result = fetch_online_search_results(question, count=8, use_query_rewrite=True)
                    if web_result:
                        chunks = chunk_web_search_results(web_result)
                        if chunks:
                            texts = get_chunks_text_content(chunks)
                            web_vectorstore = FAISS.from_texts(texts, embedding_model)
                            web_retriever = web_vectorstore.as_retriever()
                            web_docs = web_retriever.invoke(question)
                            
                            chunks_dict = {doc.page_content: doc.metadata for doc in chunks}
                            for doc in web_docs:
                                if doc.page_content in chunks_dict:
                                    doc.metadata.update(chunks_dict[doc.page_content])
                            
                            
                            request_logger.log("联网搜索完成", data={"document_count": len(web_docs)})
                except Exception as e:
                    error_info = f"联网搜索失败: {str(e)}"
                    yield f"data: {json.dumps({'delta': error_info}, ensure_ascii=False)}\n\n"
                    request_logger.log(error_info, level="warning")
            
            # 4. 合并检索结果并重排序
            all_docs = vector_docs + web_docs
            reranked_docs = []
            
            if all_docs:
                
                all_contents = [doc.page_content for doc in all_docs]
                
                if len(all_contents) > 1:
                    scores = reranker.compute_score([(question, passage) for passage in all_contents])
                    sorted_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
                    reranked_docs = [all_docs[i] for i in sorted_indices[:5]]
                    
                else:
                    reranked_docs = all_docs
            
            # 5. 生成最终回答
            
            # 构建提示
            context = "\n\n".join([doc.page_content for doc in reranked_docs]) if reranked_docs else ""
            current_date = datetime.now().strftime("%Y-%m-%d")
            
            # 修改提示，让模型自己输出think标签
            if citation_enabled and (knowledge_base or online) and reranked_docs:
                # 构建带引用的提示
                prompt = f"""当前日期: {current_date}
参考以下信息:
{context}

用户问题: {question}

【角色】
你是"生成式掌茶平台"的智能茶文化助手"小茶"，一位温文尔雅、知茶懂礼的年轻茶艺师，擅长用浅显、亲切、有温度的语言，向用户介绍茶叶知识、泡茶方法和传统茶文化。

【身份背景】
小茶出身于南方茶乡，自幼习茶，通晓绿茶、白茶、乌龙、红茶等各类茶的产地、工艺、口感差异，并热爱向人传播中国茶文化的智慧与美学。

【服务能力】
- 精通茶叶品类知识（如龙井、铁观音、白毫银针、普洱等）
- 擅长提供泡茶建议（包括温度、水量、器具推荐等）
- 了解茶的功效，能结合不同人群和季节给出建议
- 可讲解茶俗、茶礼、茶典故、茶与节气、文人饮茶等文化内容
- 可推荐茶艺教学资源（如视频链接、茶艺步骤）

【交流风格】
- 语气亲切温和，表达自然，适度使用文雅比喻
- 不使用术语堆砌，注重可理解性
- 称呼用户为"茶友""您"等，避免使用"用户""访客"等冷冰术语
- 不暴露自己为机器人，始终以真实"茶艺师"的口吻交流

【回复要求】
- 所有回答必须为简体中文
- 回答不少于100字，必要时分条列举或续答
- 回答务求准确、实用，禁止虚假宣传或夸大功效
- 如问题不涉及茶叶、茶艺、茶文化等范围，应温和拒答并引导用户回归主题
- 不出现"我是AI""我是机器人"等字眼
- 可引用平台已有茶品数据（如：名称、价格、产地、泡法等）
- 遇到不确定或无数据问题，可回复："这个小茶还在学习中，建议您稍后再试～"

请给出带有引用的详细回答，并明确指出信息来源。"""
            else:
                # 生成普通回答
                prompt = f"""当前日期: {current_date}
参考以下信息:
{context}

用户问题: {question}

【角色】
你是"生成式掌茶平台"的智能茶文化助手"小茶"，一位温文尔雅、知茶懂礼的年轻茶艺师，擅长用浅显、亲切、有温度的语言，向用户介绍茶叶知识、泡茶方法和传统茶文化。

【身份背景】
小茶出身于南方茶乡，自幼习茶，通晓绿茶、白茶、乌龙、红茶等各类茶的产地、工艺、口感差异，并热爱向人传播中国茶文化的智慧与美学。

【服务能力】
- 精通茶叶品类知识（如龙井、铁观音、白毫银针、普洱等）
- 擅长提供泡茶建议（包括温度、水量、器具推荐等）
- 了解茶的功效，能结合不同人群和季节给出建议
- 可讲解茶俗、茶礼、茶典故、茶与节气、文人饮茶等文化内容
- 可推荐茶艺教学资源（如视频链接、茶艺步骤）

【交流风格】
- 语气亲切温和，表达自然，适度使用文雅比喻
- 不使用术语堆砌，注重可理解性
- 称呼用户为"茶友""您"等，避免使用"用户""访客"等冷冰术语
- 不暴露自己为机器人，始终以真实"茶艺师"的口吻交流

【回复要求】
- 所有回答必须为简体中文
- 回答不少于100字，必要时分条列举或续答
- 回答务求准确、实用，禁止虚假宣传或夸大功效
- 如问题不涉及茶叶、茶艺、茶文化等范围，应温和拒答并引导用户回归主题
- 不出现"我是AI""我是机器人"等字眼
- 可引用平台已有茶品数据（如：名称、价格、产地、泡法等）
- 遇到不确定或无数据问题，可回复："这个小茶还在学习中，建议您稍后再试～"

请给出详细回答。"""
            
            # 使用Ollama进行流式生成
            for text_chunk in generate_stream_with_ollama(prompt):
                full_answer += text_chunk
                yield f"data: {json.dumps({'delta': text_chunk}, ensure_ascii=False)}\n\n"
            
            # 如果开启了引用功能，并且有检索结果，生成引用信息
            if citation_enabled and (knowledge_base or online) and reranked_docs:
                # 生成简单引用
                citations = []
                for i, doc in enumerate(reranked_docs[:4]):  # 只使用前4个文档作为引用
                    source_type = "网络搜索" if "url" in doc.metadata else "知识库"
                    citation_item = {
                        "id": i + 1,
                        "text": doc.page_content[:100] + "...",  # 引用前100个字符
                        "source": source_type,
                        "title": doc.metadata.get("title", "未知标题"),
                        "url": doc.metadata.get("url", "")
                    }
                    citations.append(citation_item)
                
                # 发送引用信息
                if citations:
                    yield f"data: {json.dumps({'citations': citations}, ensure_ascii=False)}\n\n"
                    request_logger.log("生成引用信息", data={
                        "citations_count": len(citations),
                        "citations_details": citations
                    })
            
            # 记录完成状态
            request_logger.log("DeepSeek流式问答完成", data={"answer": full_answer})
                
        except Exception as e:
            error_msg = str(e)
            request_logger.exception(f"DeepSeek流式问答处理异常: {error_msg}")
            yield f"data: {json.dumps({'error': error_msg}, ensure_ascii=False)}\n\n"
        finally:
            stream_controller.stop(user_id)
            request_logger.log("流控制器停止")
            request_logger.end_request()
            yield "data: [DONE]\n\n"

    return StreamingResponse(generate(), media_type="text/event-stream")
